# 加载图像数据和增强
from keras.preprocessing.image import ImageDataGenerator

# 灰度归一化
train_datagen = ImageDataGenerator(rescale=1. / 255)
# 加载数据
training_set = train_datagen.flow_from_directory('./dataset/training_set', target_size=(50, 50), batch_size=32,
                                                 class_mode='binary')
# 建立CNN模型
from keras.models import Sequential
from keras.layers import Conv2D, MaxPooling2D, Flatten, Dense

model = Sequential()
# 卷积层
model.add(Conv2D(32, (3, 3), activation='relu', input_shape=(50, 50, 3)))
# 池化层
model.add(MaxPooling2D(pool_size=(2, 2)))
# 卷积层
model.add(Conv2D(32, (3, 3), activation='relu', input_shape=(50, 50, 3)))
# 池化层
model.add(MaxPooling2D(pool_size=(2, 2)))
# 展开层
model.add(Flatten())
# MLP
model.add(Dense(units=128, activation='relu'))
model.add(Dense(units=1, activation='sigmoid'))
# 配置模型
model.compile(optimizer='adam', loss='binary_crossentropy', metrics=['accuracy'])
# 查看模型结构
model.summary()

# 将模型进行持久化
# filepath = "./1.hdf5"
# from keras.callbacks import ModelCheckpoint
# checkpoint = ModelCheckpoint(filepath, monitor='val_acc', verbose=1, save_best_only=True, mode='max')
# callbacks_list = [checkpoint]

# 训练模型
model.fit_generator(training_set, epochs=5, verbose=1,
                    # callbacks=callbacks_list
                    )

# 评估模型
# 基于训练数据的准确率
accuracy_train = model.evaluate_generator(training_set)
print('accuracy_train', accuracy_train)
# 测试数据的准确率
# 灰度归一化
test_datagen = ImageDataGenerator(rescale=1. / 255)
# 加载数据
test_set = test_datagen.flow_from_directory('./dataset/test_set', target_size=(50, 50), batch_size=32,
                                            class_mode='binary')
accuracy_test = model.evaluate_generator(test_set)
print('accuracy_test', accuracy_test)

# 识别自己的狗图
# 加载单张图片
from keras.preprocessing.image import load_img, img_to_array

pic_dog = 'pic_dog.jpeg'
pic_dog = load_img(pic_dog, target_size=(50, 50))
pic_dog = img_to_array(pic_dog)
pic_dog = pic_dog / 255
pic_dog = pic_dog.reshape(1, 50, 50, 3)
# import numpy as np

result = model.predict(pic_dog)
print('result_dog：', result)

pic_cat = 'pic_cat.jpeg'
pic_cat = load_img(pic_cat, target_size=(50, 50))
pic_cat = img_to_array(pic_cat)
pic_cat = pic_cat / 255
pic_cat = pic_cat.reshape(1, 50, 50, 3)

result = model.predict(pic_cat)
print('result_cat：', result)

# result_dog： [[0.98801017]]
# result_cat： [[0.04002476]]

# result_dog： [[0.38047832]]
# result_cat： [[0.03527132]]

# result_dog： [[0.97920275]]
# result_cat： [[8.252981e-08]]


print(training_set.class_indices)
